package org.skipper.spring.cloud.dubbo.helper;

import lombok.extern.log4j.Log4j2;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.*;
import org.springframework.util.StringUtils;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;

import static org.apache.dubbo.common.constants.CommonConstants.CONSUMER;
import static org.apache.dubbo.common.constants.CommonConstants.PROVIDER;

/**
 * Created with IntelliJ IDEA.
 * Description:
 *
 * @author shoujian
 * Date: 2022-02-14
 * Time: 11:42 AM
 */
@Activate(group = {CONSUMER,PROVIDER})
@Log4j2
public class SpringCloudDubboContext implements Filter {

    private static final String TRAFFIC_SOURCES = "trafficSources";

    private static final String TENANT_ID = "tenantId";

    private static final String TEENAGER = "teenager";

    private static final String GREY_TAG = "grey";

    @Override
    public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {

        HttpServletRequest request = ((ServletRequestAttributes) (RequestContextHolder.currentRequestAttributes())).getRequest();

        RpcContext context = RpcContext.getContext();
        //插入流量来源、租户id、青少年标识等
        context.setInvoker(invoker)
                .setInvocation(invocation)
                .setLocalAddress(invoker.getUrl().getHost(), invoker.getUrl().getPort())
                .setAttachment("trafficSources", request.getHeader(TRAFFIC_SOURCES))
                .setAttachment("tenantId", request.getHeader(TENANT_ID))
                .setAttachment("teenager", request.getHeader(TEENAGER))
        ;

        String greyTag = request.getHeader(GREY_TAG);

        if (!StringUtils.isEmpty(greyTag)) {
            context.setAttachment(Constants.FORCE_USE_TAG, greyTag);
            context.setAttachment("force", false);
        } else  {
            context.setAttachment(Constants.FORCE_USE_TAG, "pro");
            context.setAttachment("force", true);
        }

        if (invocation instanceof RpcInvocation) {
            ((RpcInvocation) invocation).setInvoker(invoker);
        }

        try {
            context.clearAfterEachInvoke(false);
            return invoker.invoke(invocation);
        } finally {
            context.clearAfterEachInvoke(true);
            // IMPORTANT! For async scenario, we must remove context from current thread, so we always create a new RpcContext for the next invoke for the same thread.
            RpcContext.removeContext(true);
            RpcContext.removeServerContext();
        }

    }
}
